#include <sstream>
#include <string>
using namespace std;
#include "clang/AST/Decl.h"
#include "clang/AST/DeclGroup.h"
#include "clang/AST/DeclarationName.h"
#include "clang/AST/StmtVisitor.h"
#include "clang/AST/DeclVisitor.h"
#include "clang/AST/TypeLocVisitor.h"
#include "clang/AST/AST.h"
#include "clang/AST/ASTConsumer.h"
#include "clang/AST/ParentMap.h"
#include "clang/Basic/SourceManager.h"
#include "clang/Basic/IdentifierTable.h"
// #include "clang/Basic/Diagnostic.h"
#include "clang/Frontend/DiagnosticOptions.h"
#include "clang/Frontend/TextDiagnosticPrinter.h"
#include "clang/Index/DeclReferenceMap.h"
#include "clang/Lex/Lexer.h"
#include "clang/Rewrite/ASTConsumers.h"
#include "clang/Rewrite/Rewriter.h"
#include "llvm/Support/MemoryBuffer.h"
#include "llvm/Support/raw_ostream.h"
#include "llvm/ADT/StringExtras.h"
#include "llvm/ADT/SmallPtrSet.h"
#include "llvm/ADT/OwningPtr.h"
#include "llvm/ADT/DenseSet.h"

#include "LoopInfo.h"
#include "SubscriptExpr.h"
#include "LoopAnalyzer.h"

using namespace clang;
using namespace idx;
    
LoopAnalyzer::LoopAnalyzer(std::string& inFile, // llvm::raw_ostream *OS,
                           Diagnostic &D, const LangOptions &LOpts)
    : Diags(D), LangOpts(LOpts), InFileName(inFile) // , OutFile(OS)
{
    CurrentDecl = 0;
}
    
LoopAnalyzer::~LoopAnalyzer() {}

void LoopAnalyzer::Initialize(ASTContext &context) {
    Context = &context;
    SM = &Context->getSourceManager();
    TUDecl = Context->getTranslationUnitDecl();
    // Get the ID and start/end of the main file.
    MainFileID = SM->getMainFileID();
    const llvm::MemoryBuffer *MainBuf = SM->getBuffer(MainFileID);
    MainFileStart = MainBuf->getBufferStart();
    MainFileEnd = MainBuf->getBufferEnd();
    Rewrite.setSourceMgr(Context->getSourceManager(), Context->getLangOptions());
    loopInfo.setRewriter(Rewrite, MainFileID);
}

// error handling interface
void LoopAnalyzer::reportError(char* msg) {
    llvm::errs() << msg;
    exit(1);
    /* Diags.Report(Diags.getCustomDiagID
       (Diagnostic::Error,
       "condition var is not same as index var\n")); */
}
    
// Process a for loop, solving dependency equations and emitting the comment.
void LoopAnalyzer::processForLoop(ForStmt* fstmt) {
    InductionVarInfo ivarInfo; // struct to hold variable and bounds.

    // PART 1: Process the initializer.
    const Stmt* initstmt = fstmt->getInit();
    if (! (initstmt && BinaryOperator::classof(initstmt))) {
        reportError(" Error: incorrect or no loop initializer found\n");
    }
      
    const BinaryOperator* initop = (const BinaryOperator*) initstmt;
    if (! initop->isAssignmentOp()) {
        reportError(" Error: initializer is not an assignment\n");
    }
    // Get variable index.
    if (! (DeclRefExpr::classof(initop->getLHS()))) {
        reportError(" Error: unrecognized LHS for loop initializer\n");
    }
    ivarInfo.name = ((DeclRefExpr*) initop->getLHS())->getDecl()->getNameAsString();
    llvm::errs() << " Index Variable: " << ivarInfo.name << "\n";

    //initlhs->dump(llvm::outs(), *SM);
    // Get value of initializer (should be constant)
    Expr* initval = initop->getRHS();
    SourceLocation initloc = initval->getExprLoc();
    if (initval->isIntegerConstantExpr(ivarInfo.lbound, *Context, &initloc)) {
        llvm::errs() << " Lower bound: " << ivarInfo.lbound << "\n";
    }
    else reportError(" Error: lower bound is not integer const expression\n");
      
    // ---- PART 2: Condition ----
    Expr* condexpr = fstmt->getCond();
    /* llvm::outs() << "dump of condition expression:\n";
       condexpr->dump(llvm::errs(), *SM); */
    if (! (BinaryOperator::classof(condexpr))) {
        reportError(" Error: condition is not binary operator\n");
    }
    // see which side is the variable
    BinaryOperator* condop = (BinaryOperator*) condexpr;
    Expr* condvarexpr;
    Expr* boundexpr;
    if (ImplicitCastExpr::classof(condop->getRHS())) { // RHS is variable.
        reportError("RHS is variable. Sorry, not supported yet.\n");
        condvarexpr = ((ImplicitCastExpr*) condop->getRHS())->getSubExpr();
        boundexpr = condop->getLHS();
    }
    else if (ImplicitCastExpr::classof(condop->getLHS())) { // LHS is variable.
        // llvm::outs() << "LHS is variable.\n";
        condvarexpr = ((ImplicitCastExpr*) condop->getLHS())->getSubExpr();
        boundexpr = condop->getRHS();
    }
    else reportError("Error: could not find variable expression in condition\n");
      
    // Make sure it corresponds to the index variable.
    if (! (DeclRefExpr::classof(condvarexpr)))  //probably never get here.
        reportError("Error: condition expression not a variable.\n");
    if (((DeclRefExpr*) condvarexpr)->getDecl()->getNameAsString() != ivarInfo.name) 
        reportError("Error: condition var is not same as index var\n");
      
    //  get operation to set ubound (for now, < or <=)
    SourceLocation exprloc = boundexpr->getExprLoc();
    if (boundexpr->isIntegerConstantExpr(ivarInfo.ubound, *Context, &exprloc)){
        // don't forget to check the op.
        string compareop = BinaryOperator::getOpcodeStr(condop->getOpcode());
        llvm::outs() << " Operator: " << compareop;
        if (compareop == "<") { ivarInfo.ubound--; }
        else if (compareop != "<=") 
            reportError("Error: unrecognized comparison operator\n");
        llvm::outs() << " Upper bound: " << ivarInfo.ubound << "\n";
    }
    else reportError(" Error: upper bound is not integer const expression\n");
      
    // PART 3: Increment 
    Expr* incexpr = fstmt->getInc();
    /* llvm::outs() << "dump of increment expression:\n";
       incexpr->dump(llvm::outs(), *SM); */
    if (UnaryOperator::classof(incexpr)) { // ++
        UnaryOperator* incop = (UnaryOperator*) incexpr;
        if (UnaryOperator::getOpcodeStr(incop->getOpcode()) != "++") {
            reportError(" Sorry, only ++ is currently supported.\n");
        }
        if (! (DeclRefExpr::classof(incop->getSubExpr()))) { //probably never.
            reportError(" Error: unrecognized increment type\n");
        }
        // check that it's the index variable being compared to a constant.
        if ( ((DeclRefExpr*) incop->getSubExpr())->getDecl()->getNameAsString()
             != ivarInfo.name ) {
            reportError(" Error: loop increment is not to index var \n");
        }
    }
    ivarInfo.step = 1;

    // Add all expressions to the loop info.
    loopInfo.push_back(ivarInfo);
    llvm::outs() << " Inserted induction var " << ivarInfo.name << " at position "
                 << loopInfo.getIndex(ivarInfo.name) << " with bounds " 
                 << ivarInfo.lbound << " and " << ivarInfo.ubound << "\n";
      
    // ok then, try the body
    Stmt* loopBody = fstmt->getBody();
    // llvm::outs() << "dump of body expression:\n";
    // bodystmt->dump(llvm::outs(), *SM);
    llvm::outs() << "\n\n"; // some space for readability.

    // If it's a nested for, recurse. 
    if (ForStmt::classof(loopBody)) {
        llvm::outs() << "Found nested for loop.\n";
        processForLoop((ForStmt*) loopBody);
        return; 
    }
      
    // 4. Process the loop body to find reads and writes. 
    vector<SubscriptExpr> reads;
    vector<SubscriptExpr> writes;
    llvm::outs() << "Processing loop body...\n";
      
    // determine whether a single or compound statement.
    int stmtIndex = 0;
    bool multiStatement = false;
    Stmt** bodyIter = &loopBody; // default case, no compound statement
    if (CompoundStmt::classof(loopBody)) {
        multiStatement = true;
        bodyIter = ((CompoundStmt*) loopBody)->body_begin();
    }

    do { // loop over all statements
        Stmt* bodystmt = *bodyIter;
        // add the source location of the statement, for inserting comments later.
        loopInfo.addStmtLoc(bodystmt->getLocStart());
        loopInfo.addStmt(bodystmt);

        // make sure it's an assignment.
        if (! (BinaryOperator::classof(bodystmt) &&
               BinaryOperator::getOpcodeStr(((BinaryOperator*) bodystmt)->getOpcode()) == "=")) {
            reportError("Error: body statement is not array assignment.\n");
        }
          
        // LHS should be ArraySubscriptExpr, (get the array variable).
        Expr* lhsExpr = ((BinaryOperator*) bodystmt)->getLHS();
        if (! (ArraySubscriptExpr::classof(lhsExpr))) 
            reportError("Error: LHS not an array subscript expression\n");
      
        // let's just assume the array base is an implicitCastExpr followed by a DeclRefExpr.
        Expr* arrayBase = ((ArraySubscriptExpr*) lhsExpr)->getBase();
        DeclRefExpr* arrayvar = (DeclRefExpr*) (((ImplicitCastExpr*) arrayBase)->getSubExpr());
        string arrayname = arrayvar->getDecl()->getNameAsString();
        // llvm::outs() << "Array name is: " << arrayname << "\n";
      
        // parse the subscript expression and add it to the 'write's
        Expr* idxExpr = ((ArraySubscriptExpr*) lhsExpr)->getIdx();
        // llvm::outs() << "\n Dump of LHS subscript expression\n";
        // idxExpr->dump(llvm::outs(), *SM);
        SubscriptExpr lhsSubscript = parseSubscriptExpr(idxExpr, arrayname, stmtIndex);
        writes.push_back(lhsSubscript);
        llvm::outs() << "lhs subscript expr: " << lhsSubscript << "\n";
          
        Expr* rhsexpr = ((BinaryOperator*) bodystmt)->getRHS();
        // llvm::outs() << "The RHS expression: ";
        // rhsexpr->dump(llvm::outs(), *SM);
        vector<SubscriptExpr> newReads = parseArrayReads(rhsexpr, stmtIndex);
        reads.insert(reads.end(), newReads.begin(), newReads.end()); 

        bodyIter++;
        stmtIndex++;
    } while (multiStatement && bodyIter != ((CompoundStmt*)loopBody)->body_end());
               
    // debug output.
    llvm::outs() << reads.size() << " Array reads found. coefficient lists: \n";
    for (int i=0; i < reads.size(); i++) {
        llvm::outs() << reads[i] << "\n";
    }

    // copy reads and writes back into loopInfo object.
    loopInfo.reads = reads;
    loopInfo.writes = writes;
      
}

/** Find subexpressions referring to the given array and read their expressions. */
vector<SubscriptExpr> LoopAnalyzer::parseArrayReads(Expr* expr, int stmtIndex) {
    vector<SubscriptExpr> result;
    // RHS may be ArraySubscriptExpr or BinaryOperator (multiple ASE's)
    if (BinaryOperator::classof(expr)) {
        result = parseArrayReads(((BinaryOperator*) expr)->getLHS(), stmtIndex);
        vector<SubscriptExpr> res2 = parseArrayReads(((BinaryOperator*) expr)->getRHS(),
                                                     stmtIndex);
        result.insert(result.end(), res2.begin(), res2.end());
    }
    else if (ImplicitCastExpr::classof(expr)) {
        // important case: check the array name and parse the subscript expression.
        if (ArraySubscriptExpr::classof(((ImplicitCastExpr*) expr)->getSubExpr())) {
            ArraySubscriptExpr* arrayExp =
                (ArraySubscriptExpr*) ((ImplicitCastExpr*) expr)->getSubExpr();
            // check that it's the right array name.
            DeclRefExpr* arrayvar =
                (DeclRefExpr*) (((ImplicitCastExpr*) arrayExp->getBase())->getSubExpr());
            std::string arrayname = arrayvar->getDecl()->getNameAsString();
            // pass the subscript itself to parseSubscriptExpr
            // llvm::outs() << "Read from array " << arrayname << " found!\n";
            result.push_back(parseSubscriptExpr(arrayExp->getIdx(), arrayname, stmtIndex));
            return result;                                   
        } else {
            // llvm::errs() << " Unknown array subscript type\n";
            return result;
        }
    }
    else if (IntegerLiteral::classof(expr) || FloatingLiteral::classof(expr)) {
        return result;
    }
    else { llvm::errs() << " unrecognized expression type in RHS\n"; return result; }

}

/** Parse an expression representing a single position of a subscript. */
SubscriptExpr LoopAnalyzer::parseSubscriptExpr(Expr* expr, std::string arrayname, int stmtIndex) {
    // initialize coefficients of the linear subscript expression
    SubscriptExpr result(loopInfo.size(), arrayname, stmtIndex);
      
    if (ImplicitCastExpr::classof(expr))
        expr = ((ImplicitCastExpr*)expr)->getSubExpr();
    // base case - a single variable reference.
    if (DeclRefExpr::classof(expr)) {
        string varname = ((DeclRefExpr*) expr)->getDecl()->getNameAsString();
        int varidx = loopInfo.getIndex(varname);
        if (varidx != -1) {
            // llvm::outs() << "Prev value: " << result[varidx] << "\n";
            result[varidx]++;
            //llvm::outs() << "Incrementing subscr ref to " << varname
            //             << ", index " << varidx << " to " << result[varidx] << "\n";
        }
    }
    else if (BinaryOperator::classof(expr)) {
        // i - 1, for example
        string optype = BinaryOperator::getOpcodeStr(((BinaryOperator*) expr)->getOpcode());
        SubscriptExpr res_l = parseSubscriptExpr(((BinaryOperator*) expr)->getLHS(),
                                                 arrayname, stmtIndex);
        SubscriptExpr res_r = parseSubscriptExpr(((BinaryOperator*) expr)->getRHS(),
                                                 arrayname, stmtIndex);
        if (optype == "+") {
            result +=  res_l + res_r;
        }
        else if (optype == "-") {
            result += res_l - res_r;
        }
        else {
            // Constant multiplication not supported.
            llvm::errs() << "Error: unhandled operator " << optype << " in subscript\n";
        }
    }
    else if (IntegerLiteral::classof(expr)) {
        result[0] = ((IntegerLiteral*) expr)->getValue();
        // llvm::outs() << "Integer literal subscript found.\n";
    }
        
    return result;
}
  
// This is called after parsing; it gets the declarations of the TU and VISITs each one.
void LoopAnalyzer::HandleTranslationUnit(ASTContext &C) {
    // llvm::outs() << "Handling TU.\n";
    // modify calls
    TranslationUnitDecl *TU = Context->getTranslationUnitDecl();
    for (DeclContext::decl_iterator I = TU->decls_begin(),
             E = TU->decls_end(); I != E; ++I) {
        Decl *D = *I;  //  I is an iterator, so a pointer to a pointer.
        Visit(D); // What type is D? get some debug output...
    }
    // Get the buffer corresponding to MainFileID.
    // If we haven't changed it, then we are done.
    // MOVED to LoopInfo.
}
  
string* LoopAnalyzer::getCodeString(SourceRange Range) {
    return new std::string(Context->getSourceManager()
                           .getCharacterData(Range.getBegin()),
                           Rewrite.getRangeSize(Range));
}

//-------------------------------------------------------------
// It's all boilerplate below here. Please ignore :)
//-------------------------------------------------------------
void LoopAnalyzer::VisitDeclRefExpr(DeclRefExpr* E) { // superclass of Expr and Stmt.
}

void LoopAnalyzer::Visit(Decl *D) {
    if (!D) return;
    Decl *PrevDecl = CurrentDecl;
    CurrentDecl = D;
    // The visitor pattern callback that causes correct method to be dispatched?
    BaseDeclVisitor::Visit(D); 
    CurrentDecl = PrevDecl;
}

void LoopAnalyzer::VisitDeclaratorDecl(DeclaratorDecl *D) {
    BaseDeclVisitor::VisitDeclaratorDecl(D);
    if (TypeSourceInfo *TInfo = D->getTypeSourceInfo())
        Visit(TInfo->getTypeLoc());
}

void LoopAnalyzer::VisitFunctionDecl(FunctionDecl *D) {
    BaseDeclVisitor::VisitFunctionDecl(D);
    if (D->isThisDeclarationADefinition()) {
        Visit(D->getBody());
    }
}

void LoopAnalyzer::VisitObjCMethodDecl(ObjCMethodDecl *D) {
    BaseDeclVisitor::VisitObjCMethodDecl(D);
    if (D->getBody())
        Visit(D->getBody());
}

void LoopAnalyzer::VisitBlockDecl(BlockDecl *D) {
    BaseDeclVisitor::VisitBlockDecl(D);
    Visit(D->getBody());
}

void LoopAnalyzer::VisitVarDecl(VarDecl *D) {
    BaseDeclVisitor::VisitVarDecl(D);
    // llvm::errs() << "Visited a VarDecl!\n";
    if (Expr *Init = D->getInit()) // Is variable initialized at declaration?
        Visit(Init);
}

void LoopAnalyzer::VisitDecl(Decl *D) {
    if (isa<FunctionDecl>(D) || isa<ObjCMethodDecl>(D) || isa<BlockDecl>(D))
        return;
    if (DeclContext *DC = dyn_cast<DeclContext>(D))
        static_cast<LoopAnalyzer*>(this)->VisitDeclContext(DC);
}

void LoopAnalyzer::VisitDeclContext(DeclContext *DC) {
    for (DeclContext::decl_iterator
             I = DC->decls_begin(), E = DC->decls_end(); I != E; ++I)
        Visit(*I);
}

void LoopAnalyzer::VisitDeclStmt(DeclStmt *Node) {
    // llvm::errs() << "Visiting DeclStmt.\n";
    for (DeclStmt::decl_iterator
             I = Node->decl_begin(), E = Node->decl_end(); I != E; ++I)
        Visit(*I);
}

void LoopAnalyzer::VisitBlockExpr(BlockExpr *Node) {
    // The BlockDecl is also visited by 'VisitDeclContext()'.  No need to visit it twice.
}

void LoopAnalyzer::VisitStmt(Stmt *Node) {

    // check for "for" statement
    if (isa<ForStmt>(Node)) {
        ForStmt *fstmt = (ForStmt*) Node;
        llvm::outs() << "Found an outer FOR loop!\n";
        // create empty loop nest and call the handler method.
        // LoopNestInfo loopinfo; 
        processForLoop(fstmt);

    } // isa<ForStmt>
    else {
            
        for (Stmt::child_iterator
                 I = Node->child_begin(), E = Node->child_end(); I != E; ++I)
            if (*I)
                Visit(*I);
    }
}

void LoopAnalyzer::Visit(TypeLoc TL) {
    for (; TL; TL = TL.getNextTypeLoc())
        BaseTypeLocVisitor::Visit(TL);
}

void LoopAnalyzer::VisitArrayLoc(ArrayTypeLoc TL) {
    BaseTypeLocVisitor::VisitArrayTypeLoc(TL); // just calls base class
    if (TL.getSizeExpr())
        Visit(TL.getSizeExpr());
    llvm::errs() << "arrayLoc found";
}

void LoopAnalyzer::VisitFunctionTypeLoc(FunctionTypeLoc TL) {
    BaseTypeLocVisitor::VisitFunctionTypeLoc(TL);
    for (unsigned i = 0; i != TL.getNumArgs(); ++i)
        Visit(TL.getArg(i));
}

